package onnxruntime

import (
	"fmt"
	"unsafe"
)

// #cgo CFLAGS: -O2 -g
//
// #include "onnxruntime_wrapper.h"
import "C"

// This wraps internal implementation details to avoid exposing them to users
// via the ArbitraryTensor interface.
type TensorInternalData struct {
	OrtValue *C.OrtValue
}

// An interface for managing tensors where we don't care about accessing the
// underlying data slice. All typed tensors will support this interface,
// regardless of the underlying data type.
type ArbitraryTensor interface {
	DataType() C.ONNXTensorElementDataType
	GetShape() Shape
	Destroy() error
	GetInternals() *TensorInternalData
	ZeroContents()
}

// Used to manage all input and output data for onnxruntime networks. A Tensor
// always has an associated type and refers to data contained in an underlying
// Go slice. New tensors should be created using the NewTensor or
// NewEmptyTensor functions, and must be destroyed using the Destroy function
// when no longer needed.
type Tensor[T TensorData] struct {
	// The shape of the tensor
	shape Shape
	// The go slice containing the flattened data that backs the ONNX tensor.
	data []T
	// The number of bytes taken by the data slice.
	dataSize uintptr
	// The underlying ONNX value we use with the C API.
	ortValue *C.OrtValue
}

// Cleans up and frees the memory associated with this tensor.
func (t *Tensor[_]) Destroy() error {
	C.ReleaseOrtValue(t.ortValue)
	t.ortValue = nil
	t.data = nil
	t.dataSize = 0
	t.shape = nil
	return nil
}

// Returns the slice containing the tensor's underlying data. The contents of
// the slice can be read or written to get or set the tensor's contents.
func (t *Tensor[T]) GetData() []T {
	return t.data
}

// Returns the value from the ONNXTensorElementDataType C enum corresponding to
// the type of data held by this tensor.
//
// NOTE: This function was added prior to the introduction of the
// Go TensorElementDataType int wrapping the C enum, so it still returns the
// CGo type.
func (t *Tensor[T]) DataType() C.ONNXTensorElementDataType {
	return GetTensorElementDataType[T]()
}

// Returns the shape of the tensor. The returned shape is only a copy;
// modifying this does *not* change the shape of the underlying tensor.
// (Modifying the tensor's shape can only be accomplished by Destroying and
// recreating the tensor with the same data.)
func (t *Tensor[_]) GetShape() Shape {
	return t.shape.Clone()
}

func (t *Tensor[_]) GetInternals() *TensorInternalData {
	return &TensorInternalData{
		OrtValue: t.ortValue,
	}
}

// Sets every element in the tensor's underlying data slice to 0.
func (t *Tensor[T]) ZeroContents() {
	C.memset(unsafe.Pointer(&t.data[0]), 0, C.size_t(t.dataSize))
}

// Makes a deep copy of the tensor, including its ONNXRuntime value. The Tensor
// returned by this function must be destroyed when no longer needed. The
// returned tensor will also no longer refer to the same underlying data; use
// GetData() to obtain the new underlying slice.
func (t *Tensor[T]) Clone() (*Tensor[T], error) {
	toReturn, e := NewEmptyTensor[T](t.shape)
	if e != nil {
		return nil, fmt.Errorf("error allocating tensor clone: %w", e)
	}
	copy(toReturn.GetData(), t.data)
	return toReturn, nil
}

// Creates a new empty tensor with the given shape. The shape provided to this
// function is copied, and is no longer needed after this function returns.
func NewEmptyTensor[T TensorData](s Shape) (*Tensor[T], error) {
	e := s.Validate()
	if e != nil {
		return nil, fmt.Errorf("invalid tensor shape: %w", e)
	}
	elementCount := s.FlattenedSize()
	data := make([]T, elementCount)
	return NewTensor(s, data)
}

// Creates a new tensor backed by an existing data slice. The shape provided to
// this function is copied, and is no longer needed after this function
// returns. If the data slice is longer than s.FlattenedSize(), then only the
// first portion of the data will be used.
func NewTensor[T TensorData](s Shape, data []T) (*Tensor[T], error) {
	if !IsInitialized() {
		return nil, ErrorNotInitialized
	}
	e := s.Validate()
	if e != nil {
		return nil, fmt.Errorf("invalid tensor shape: %w", e)
	}
	elementCount := s.FlattenedSize()
	if elementCount > int64(len(data)) {
		return nil, fmt.Errorf("the tensor's shape (%s) requires %d elements, but only %d were provided", s, elementCount, len(data))
	}
	var ortValue *C.OrtValue
	dataType := GetTensorElementDataType[T]()
	dataSize := unsafe.Sizeof(data[0]) * uintptr(elementCount)

	status := C.CreateOrtTensorWithShape(unsafe.Pointer(&data[0]),
		C.size_t(dataSize), (*C.int64_t)(unsafe.Pointer(&s[0])),
		C.int64_t(len(s)), ortMemoryInfo, dataType, &ortValue)
	if status != nil {
		return nil, fmt.Errorf("ORT API error creating tensor: %s",
			statusToError(status))
	}

	toReturn := Tensor[T]{
		data:     data[0:elementCount],
		dataSize: dataSize,
		shape:    s.Clone(),
		ortValue: ortValue,
	}
	// TODO: Set a finalizer on new Tensors to hopefully prevent careless
	// memory leaks.
	// - Idea: use a "destroyable" interface?
	return &toReturn, nil
}
